#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Histogram Lookup Tool

This tool provides fast lookup functionality for UMA histogram names,
including reverse lookup from expanded names to original definitions.

Copyright 2024 Chromium Authors
Use of this source code is governed by a BSD-style license that can be
found in the LICENSE file.
"""

import argparse
import sqlite3
import sys
from pathlib import Path


class HistogramLookup:
    """Fast lookup for histogram information."""

    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = None

    def connect(self):
        """Connect to the database."""
        self.conn = sqlite3.connect(self.db_path)
        self.conn.row_factory = sqlite3.Row

        # Check if expanded_histograms table exists
        cursor = self.conn.cursor()
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name='expanded_histograms'
        """)
        self.has_expansion_table = cursor.fetchone() is not None
        cursor.close()

    def close(self):
        """Close database connection."""
        if self.conn:
            self.conn.close()

    def find_by_exact_name(self, name: str):
        """Find histogram by exact name (original or expanded)."""
        cursor = self.conn.cursor()

        # Try exact match in metrics table first
        cursor.execute("""
            SELECT id, name, type, summary, units, expires_after, component, file_path, line_number
            FROM metrics
            WHERE name = ? AND type = 'histogram'
        """, (name,))

        result = cursor.fetchone()
        if result:
            return {
                'id': result['id'],
                'name': result['name'],
                'type': 'original',
                'summary': result['summary'],
                'units': result['units'],
                'expires_after': result['expires_after'],
                'component': result['component'],
                'file_path': result['file_path'],
                'line_number': result['line_number']
            }

        # Try expanded names
        cursor.execute("""
            SELECT m.id, m.name, m.summary, m.units, m.expires_after, m.component,
                   m.file_path, m.line_number, eh.expanded_name, eh.token_values
            FROM expanded_histograms eh
            JOIN metrics m ON eh.original_metric_id = m.id
            WHERE eh.expanded_name = ? AND m.type = 'histogram'
        """, (name,))

        result = cursor.fetchone()
        if result:
            import json
            return {
                'id': result['id'],
                'name': result['name'],
                'type': 'expanded',
                'expanded_name': result['expanded_name'],
                'summary': result['summary'],
                'units': result['units'],
                'expires_after': result['expires_after'],
                'component': result['component'],
                'file_path': result['file_path'],
                'line_number': result['line_number'],
                'token_values': json.loads(result['token_values']) if result['token_values'] else {}
            }

        return None

    def find_by_pattern(self, pattern: str, limit: int = 20):
        """Find histograms matching a pattern."""
        cursor = self.conn.cursor()

        results = []

        # Search in original names
        cursor.execute("""
            SELECT id, name, type, summary, units, expires_after, component, file_path, line_number
            FROM metrics
            WHERE name LIKE ? AND type = 'histogram'
            ORDER BY name
            LIMIT ?
        """, (f'%{pattern}%', limit))

        for row in cursor.fetchall():
            results.append({
                'id': row['id'],
                'name': row['name'],
                'type': 'original',
                'summary': row['summary'],
                'units': row['units'],
                'expires_after': row['expires_after'],
                'component': row['component'],
                'file_path': row['file_path'],
                'line_number': row['line_number']
            })

        # Search in expanded names if table exists and we haven't reached the limit
        if self.has_expansion_table and len(results) < limit:
            remaining_limit = limit - len(results)
            cursor.execute("""
                SELECT m.id, m.name, m.summary, m.units, m.expires_after, m.component,
                       m.file_path, m.line_number, eh.expanded_name, eh.token_values
                FROM expanded_histograms eh
                JOIN metrics m ON eh.original_metric_id = m.id
                WHERE eh.expanded_name LIKE ? AND m.type = 'histogram'
                ORDER BY eh.expanded_name
                LIMIT ?
            """, (f'%{pattern}%', remaining_limit))

            for row in cursor.fetchall():
                import json
                results.append({
                    'id': row['id'],
                    'name': row['name'],
                    'type': 'expanded',
                    'expanded_name': row['expanded_name'],
                    'summary': row['summary'],
                    'units': row['units'],
                    'expires_after': row['expires_after'],
                    'component': row['component'],
                    'file_path': row['file_path'],
                    'line_number': row['line_number'],
                    'token_values': json.loads(row['token_values']) if row['token_values'] else {}
                })

        return results

    def get_expansions(self, original_name: str):
        """Get all expanded names for an original histogram."""
        if not self.has_expansion_table:
            return []

        cursor = self.conn.cursor()

        cursor.execute("""
            SELECT m.id, m.name, eh.expanded_name, eh.token_values
            FROM metrics m
            JOIN expanded_histograms eh ON eh.original_metric_id = m.id
            WHERE m.name = ? AND m.type = 'histogram'
        """, (original_name,))

        results = []
        for row in cursor.fetchall():
            import json
            results.append({
                'id': row['id'],
                'original_name': row['name'],
                'expanded_name': row['expanded_name'],
                'token_values': json.loads(row['token_values']) if row['token_values'] else {}
            })

        return results

    def search_by_token_values(self, token_values: dict):
        """Search histograms by specific token values."""
        if not self.has_expansion_table:
            return []

        cursor = self.conn.cursor()

        # Build WHERE clause for token values
        conditions = []
        params = []

        for token_key, token_value in token_values.items():
            conditions.append("eh.token_values LIKE ?")
            params.append(f'%"{token_key}": "{token_value}"%')

        if not conditions:
            return []

        where_clause = " AND ".join(conditions)

        cursor.execute(f"""
            SELECT m.name, eh.expanded_name, eh.token_values
            FROM expanded_histograms eh
            JOIN metrics m ON eh.original_metric_id = m.id
            WHERE {where_clause} AND m.type = 'histogram'
            ORDER BY m.name
        """, params)

        results = []
        for row in cursor.fetchall():
            import json
            results.append({
                'name': row['name'],
                'expanded_name': row['expanded_name'],
                'token_values': json.loads(row['token_values']) if row['token_values'] else {}
            })

        return results


def format_result(result):
    """Format lookup result for display."""
    if not result:
        return "Not found"

    output = []
    output.append(f"Name: {result['name']}")

    if result['type'] == 'expanded':
        output.append(f"Expanded: {result['expanded_name']}")
        if result.get('token_values'):
            output.append(f"Token values: {result['token_values']}")

    if result['summary']:
        output.append(f"Summary: {result['summary'][:200]}{'...' if len(result['summary']) > 200 else ''}")

    if result['units']:
        output.append(f"Units: {result['units']}")

    if result['expires_after']:
        output.append(f"Expires: {result['expires_after']}")

    if result['component']:
        output.append(f"Component: {result['component']}")

    output.append(f"Source: {Path(result['file_path']).name}:{result['line_number']}")

    return '\n  '.join(output)


def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Fast lookup for UMA histograms')
    parser.add_argument('--db', default='./uma_ukm_metrics.db',
                        help='Database file path')
    parser.add_argument('--exact',
                        help='Find histogram by exact name')
    parser.add_argument('--search',
                        help='Search histograms by pattern')
    parser.add_argument('--expansions',
                        help='Show all expansions for original histogram name')
    parser.add_argument('--token',
                        help='Search by token values (format: key=value,key2=value2)')
    parser.add_argument('--limit', type=int, default=10,
                        help='Limit number of results (default: 10)')

    args = parser.parse_args()

    if not Path(args.db).exists():
        print(f"Database file not found: {args.db}")
        print("Please run uma_ukm_extractor.py first to create the database.")
        sys.exit(1)

    lookup = HistogramLookup(args.db)
    lookup.connect()

    try:
        if args.exact:
            result = lookup.find_by_exact_name(args.exact)
            if result:
                print(f"Found histogram (type: {result['type']}):")
                print(f"  {format_result(result)}")
            else:
                print(f"Histogram not found: {args.exact}")

        elif args.search:
            results = lookup.find_by_pattern(args.search, args.limit)
            if results:
                print(f"Found {len(results)} histograms matching '{args.search}':")
                for i, result in enumerate(results, 1):
                    print(f"\n{i}. {format_result(result)}")
            else:
                print(f"No histograms found matching: {args.search}")

        elif args.expansions:
            expansions = lookup.get_expansions(args.expansions)
            if expansions:
                print(f"Found {len(expansions)} expansions for '{args.expansions}':")
                for i, expansion in enumerate(expansions, 1):
                    print(f"\n{i}. {expansion['expanded_name']}")
                    if expansion['token_values']:
                        print(f"   Tokens: {expansion['token_values']}")
            else:
                print(f"No expansions found for: {args.expansions}")

        elif args.token:
            try:
                token_dict = {}
                for pair in args.token.split(','):
                    key, value = pair.split('=', 1)
                    token_dict[key.strip()] = value.strip()

                results = lookup.search_by_token_values(token_dict)
                if results:
                    print(f"Found {len(results)} histograms with token values {token_dict}:")
                    for i, result in enumerate(results, 1):
                        print(f"\n{i}. {result['expanded_name']}")
                        print(f"   Original: {result['name']}")
                        print(f"   Tokens: {result['token_values']}")
                else:
                    print(f"No histograms found with token values: {token_dict}")

            except ValueError:
                print(f"Invalid token format. Use: key=value,key2=value2")

        else:
            print("Use --help to see available options")
            print("\nExamples:")
            print(f"  python3 {sys.argv[0]} --exact 'V8.GC.Cycle.HighPriority.CollectionRate.FullHeap'")
            print(f"  python3 {sys.argv[0]} --search 'GC.CollectionRate'")
            print(f"  python3 {sys.argv[0]} --expansions 'V8.GC.Cycle{Priority}.CollectionRate.Full{{Heap}}'")
            print(f"  python3 {sys.argv[0]} --token 'Priority=HighPriority,Heap=FullHeap'")

    finally:
        lookup.close()


if __name__ == '__main__':
    main()